#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
任务映射管理器
管理任务ID与GPU服务器的映射关系
"""

import json
import os
import time
import logging
from typing import Dict, Optional
from threading import Lock

logger = logging.getLogger(__name__)

class TaskMappingManager:
    """任务映射管理器"""
    
    def __init__(self, mapping_file: str = None):
        self.mapping_file = mapping_file or os.path.join(os.path.dirname(__file__), "task_mapping.json")
        self.task_mappings: Dict[str, Dict] = {}  # task_id -> {server_url, created_time}
        self.lock = Lock()
        self.load_mappings()
    
    def load_mappings(self):
        """加载任务映射"""
        try:
            if os.path.exists(self.mapping_file):
                with open(self.mapping_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    self.task_mappings = data.get('mappings', {})
                logger.info(f"加载了 {len(self.task_mappings)} 个任务映射")
        except Exception as e:
            logger.error(f"加载任务映射失败: {e}")
            self.task_mappings = {}
    
    def save_mappings(self):
        """保存任务映射"""
        try:
            with self.lock:
                data = {
                    'mappings': self.task_mappings,
                    'last_updated': time.time()
                }
                # 先尝试创建目录
                os.makedirs(os.path.dirname(self.mapping_file), exist_ok=True)
                with open(self.mapping_file, 'w', encoding='utf-8') as f:
                    json.dump(data, f, indent=2, ensure_ascii=False)
                logger.info(f"任务映射已保存到: {self.mapping_file}")
        except Exception as e:
            logger.error(f"保存任务映射失败: {e}")
    
    def add_task_mapping(self, task_id: str, server_url: str):
        """添加任务映射"""
        with self.lock:
            self.task_mappings[task_id] = {
                'server_url': server_url,
                'created_time': time.time()
            }
            self.save_mappings()
            logger.info(f"任务映射已添加: {task_id} -> {server_url}")
    
    def get_task_server(self, task_id: str) -> Optional[str]:
        """获取任务对应的服务器URL"""
        with self.lock:
            mapping = self.task_mappings.get(task_id)
            if mapping:
                return mapping.get('server_url')
            return None
    
    def remove_task_mapping(self, task_id: str):
        """移除任务映射"""
        with self.lock:
            if task_id in self.task_mappings:
                del self.task_mappings[task_id]
                self.save_mappings()
                logger.info(f"任务映射已移除: {task_id}")
    
    def cleanup_old_mappings(self, max_age_hours: int = 24):
        """清理过期的任务映射"""
        current_time = time.time()
        max_age_seconds = max_age_hours * 3600
        
        with self.lock:
            expired_tasks = []
            for task_id, mapping in self.task_mappings.items():
                created_time = mapping.get('created_time', 0)
                if current_time - created_time > max_age_seconds:
                    expired_tasks.append(task_id)
            
            for task_id in expired_tasks:
                del self.task_mappings[task_id]
            
            if expired_tasks:
                self.save_mappings()
                logger.info(f"清理了 {len(expired_tasks)} 个过期任务映射")
    
    def get_mapping_stats(self) -> Dict:
        """获取映射统计信息"""
        with self.lock:
            return {
                'total_mappings': len(self.task_mappings),
                'mappings': dict(self.task_mappings)
            }

# 全局实例
task_mapping_manager = TaskMappingManager()
